from cryptography.fernet import Fernet
import os
import getpass

class SecertUtil:
    def __init__(self):
        self.cur_dir = os.path.dirname(os.path.abspath(__file__))
        self.secret_path = os.path.join(self.cur_dir, 'secret')
        if not os.path.exists(self.secret_path):
            os.mkdir(self.secret_path)

    def set_secret(self):
        # 如果不存在secret.dat文件则生成该密钥文件
        if not os.path.exists(os.path.join(self.secret_path, 'secret.dat')):
            self.s_key = Fernet.generate_key()
            with open(os.path.join(self.secret_path, 'secret.dat'), mode='wb') as f:
                f.write(self.s_key)
        # 如果存在secret.dat则直接读取该密钥文件
        with open(os.path.join(self.secret_path, 'secret.dat'), mode='rb') as f:
            secret_text = f.read()
        # 返回密钥
        return secret_text

    def encrypt(self):
        key = self.set_secret()
        while True:
            username = input('请输入用户名：')
            password = getpass.getpass('请输入密码：')
            repassword = getpass.getpass('请再次确认密码：')
            if not password == repassword:
                print('两次输入密码不一致')
            else:
                break
        key_generator = Fernet(key)
        enuser = key_generator.encrypt(username.encode())
        enpwd = key_generator.encrypt(password.encode())
        return enuser, enpwd

    def decrypt(self, enuser:bytes, enpwd:bytes):
        key = self.set_secret()
        key_generator = Fernet(key)
        deuser = key_generator.decrypt(enuser)
        depwd = key_generator.decrypt(enpwd)
        return deuser.decode(), depwd.decode()
if __name__ == '__main__':
    su = SecertUtil()
    a, b = su.encrypt()
    print(a.decode() + '\n' + b.decode())
    print(su.set_secret())
    c, d = su.decrypt(a, b)
    print(c, d)

    # s_key = Fernet.generate_key()
    # key_generator = Fernet(s_key)
    # encrypted_str = key_generator.encrypt("{}".format(str_encrypted).encode())
    # print(encrypted_str)
    # decrypted_str = key_generator.decrypt(encrypted_str)
    # print(decrypted_str)

# def Decrypt(str_decrypted):
#     return key_generator.decrypt(encrypted_str)
# Encrypt('administrator')